import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Ellipse

np.random.seed(42)

# 1. Define 3 Gaussian clusters in 2D for labels y=0,1,2
means = np.array([[0, 0], [3, 3], [0, 5]])
cov = np.array([[0.3, 0], [0, 0.3]])
num_samples_per_class = 100

# Generate ground truth data
data = []
labels = []
for i, mean in enumerate(means):
    samples = np.random.multivariate_normal(mean, cov, num_samples_per_class)
    data.append(samples)
    labels.append(np.full(num_samples_per_class, i))
data = np.concatenate(data, axis=0)
labels = np.concatenate(labels, axis=0)

# Simulate "x_t" = noised samples from x_0
noise_scale = 1.0
x_t = data + np.random.randn(*data.shape) * noise_scale

# Define DINO-like embedding (simulate with a projection + noise)
A = np.random.randn(10, 2)  # projection to 10D
z_x = (A @ x_t.T).T + np.random.randn(x_t.shape[0], 10) * 0.05  # high-dim DINOv2-like embedding

# Compute class-wise anchors in embedding space
z_anchors = np.array([z_x[labels == i].mean(axis=0) for i in range(3)])

# Compute weight per anchor using softmax of squared distances in embedding space
def softmax(x):
    e_x = np.exp(-x)
    return e_x / e_x.sum(axis=1, keepdims=True)

distances = np.array([[np.linalg.norm(z - anchor)**2 for anchor in z_anchors] for z in z_x])
weights = softmax(distances)  # shape: (N, 3)

# Simulate denoising by weighted combination of anchor means
denoised_softanchor = weights @ means  # ours: weighted denoising target
denoised_baseline = means[labels]      # baseline: purely based on label

# Plot
fig, axs = plt.subplots(1, 3, figsize=(15, 5))

# Ground Truth
axs[0].scatter(data[:, 0], data[:, 1], c=labels, cmap='viridis', s=10, label="GT")
axs[0].set_title("Ground Truth Samples")
axs[0].axis('equal')

# Baseline Sampling (label only)
axs[1].scatter(x_t[:, 0], x_t[:, 1], c='gray', s=5, alpha=0.3, label='x_t')
axs[1].scatter(denoised_baseline[:, 0], denoised_baseline[:, 1], c=labels, cmap='viridis', s=10)
axs[1].set_title("Baseline: Label-Only Denoising")
axs[1].axis('equal')

# Ours: Soft Anchors (DINOv2-guided)
axs[2].scatter(x_t[:, 0], x_t[:, 1], c='gray', s=5, alpha=0.3, label='x_t')
axs[2].scatter(denoised_softanchor[:, 0], denoised_softanchor[:, 1], c=labels, cmap='viridis', s=10)
axs[2].set_title("Ours: Soft Anchor Guidance")
axs[2].axis('equal')

plt.tight_layout()



# 构造结构信息（ControlNet 样式）：模拟一条非线性边界曲线，作为结构条件
# 假设结构条件是一个弯曲路径 y = sin(x)，我们想让部分数据“贴合”这个结构进行生成

# 结构条件曲线：y = sin(x) + 偏移
curve_x = np.linspace(-1, 4, 300)
curve_y = np.sin(curve_x) + 2.5
structure_curve = np.stack([curve_x, curve_y], axis=1)

# 为 toy 数据模拟结构引导：选择离曲线最近的点作为目标样本
from scipy.spatial import cKDTree

# 构造 KD-Tree 查询每个点到曲线的最近点
tree = cKDTree(structure_curve)
_, idx = tree.query(denoised_softanchor)
guided_points = structure_curve[idx]  # 每个点被引导到的结构位置

# 加权融合结构条件与语义锚引导（Soft anchor + structure guidance）
lambda_sem = 0.7  # 语义权重
lambda_str = 1 - lambda_sem  # 结构权重
denoised_combined = lambda_sem * denoised_softanchor + lambda_str * guided_points

# Plot 三种方法对比：Baseline / Soft Anchor / Soft Anchor + Structure
fig, axs = plt.subplots(1, 3, figsize=(15, 5))

# Baseline
axs[0].scatter(x_t[:, 0], x_t[:, 1], c='gray', s=5, alpha=0.3)
axs[0].scatter(denoised_baseline[:, 0], denoised_baseline[:, 1], c=labels, cmap='viridis', s=10)
axs[0].plot(curve_x, curve_y, 'r--', label='Structure Curve')
axs[0].set_title("Baseline (Label Only)")
axs[0].legend()
axs[0].axis('equal')

# Ours: Soft Anchors
axs[1].scatter(x_t[:, 0], x_t[:, 1], c='gray', s=5, alpha=0.3)
axs[1].scatter(denoised_softanchor[:, 0], denoised_softanchor[:, 1], c=labels, cmap='viridis', s=10)
axs[1].plot(curve_x, curve_y, 'r--')
axs[1].set_title("Ours: Soft Anchor Only")
axs[1].axis('equal')

# Ours: Soft Anchors + Structure Guidance
axs[2].scatter(x_t[:, 0], x_t[:, 1], c='gray', s=5, alpha=0.3)
axs[2].scatter(denoised_combined[:, 0], denoised_combined[:, 1], c=labels, cmap='viridis', s=10)
axs[2].plot(curve_x, curve_y, 'r--')
axs[2].set_title("Ours: Soft Anchor + Structure Guidance")
axs[2].axis('equal')

plt.tight_layout()
plt.savefig('test.png')
